import streamlit as st
from PIL import Image
import numpy as np
import cv2
import os
from streamlit_drawable_canvas import st_canvas
import io
import tempfile
from work3 import shibie

st.set_page_config(
    page_title="学号识别系统",
    page_icon="✏️",
    layout="wide",
    initial_sidebar_state="expanded"
)

st.markdown("""
    <style>
    .main {
        padding: 2rem;
    }
    .stButton>button {
        width: 100%;
        border-radius: 10px;
        height: 3em;
        font-size: 16px;
    }
    .uploadedFile {
        border: 2px dashed #ccc;
        padding: 20px;
        border-radius: 10px;
    }
    .success-text {
        color: #0f5132;
        background-color: #d1e7dd;
        padding: 1rem;
        border-radius: 0.5rem;
        margin: 1rem 0;
    }
    </style>
""", unsafe_allow_html=True)


def safe_file_delete(file_path):
    """安全地删除文件，忽略任何错误"""
    try:
        if os.path.exists(file_path):
            os.unlink(file_path)
    except:
        pass


def main():
    st.title("✏️ 手写学号识别系统")
    st.markdown("---")

    col1, col2 = st.columns(2, gap="large")

    with col1:
        st.subheader("📝 手写输入")

        canvas_result = st_canvas(
            fill_color="rgba(255, 165, 0, 0.3)",
            stroke_width=10,
            stroke_color="#000000",
            background_color="#ffffff",
            width=600,
            height=200,
            drawing_mode="freedraw",
            key="canvas",
        )

        if st.button("识别手写内容", key="recognize_drawing"):
            if canvas_result.image_data is not None:
                tmp_file = None
                try:
                    with tempfile.NamedTemporaryFile(delete=False, suffix='.png') as f:
                        tmp_file = f.name
                        img_array = cv2.cvtColor(canvas_result.image_data.astype('uint8'), cv2.COLOR_RGBA2GRAY)
                        cv2.imwrite(tmp_file, img_array)
                        result = shibie(tmp_file)
                        st.markdown(f"""
                            <div class="success-text">
                                <h3>识别结果: {result}</h3>
                            </div>
                        """, unsafe_allow_html=True)
                except Exception as e:
                    st.error("识别过程中出现错误，请重试")
                finally:
                    if tmp_file:
                        safe_file_delete(tmp_file)
            else:
                st.warning("请先在画板上书写内容")

    with col2:
        st.subheader("📤 图片上传")

        uploaded_file = st.file_uploader(
            "选择包含手写学号的图片",
            type=['png', 'jpg', 'jpeg'],
            help="支持PNG、JPG格式的图片"
        )

        if uploaded_file is not None:
            image = Image.open(uploaded_file)
            st.image(image, caption="上传的图片", use_column_width=True)

            if st.button("识别上传的图片", key="recognize_upload"):
                tmp_file = None
                try:
                    with tempfile.NamedTemporaryFile(delete=False, suffix='.png') as f:
                        tmp_file = f.name
                        image.save(tmp_file, 'PNG')
                        result = shibie(tmp_file)
                        st.markdown(f"""
                            <div class="success-text">
                                <h3>识别结果: {result}</h3>
                            </div>
                        """, unsafe_allow_html=True)
                except Exception as e:
                    st.error("识别过程中出现错误，请重试")
                finally:
                    if tmp_file:
                        safe_file_delete(tmp_file)

    st.markdown("---")
    st.markdown("""
    ### 📌 使用说明
    1. **手写输入**: 在左侧画板上直接书写学号
    2. **图片上传**: 上传含有手写学号的图片文件
    3. **注意事项**: 
        - 请确保字迹清晰
        - 数字之间留有适当间距
    """)


if __name__ == "__main__":
    main()